use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use dashmap::DashMap;
use datafusion::arrow::datatypes::{Schema, SchemaRef};
use datafusion::execution::SendableRecordBatchStream;
use datafusion::physical_plan::memory::MemoryStream;
use datafusion::prelude::DataFrame;
use futures_util::StreamExt;

use common::{Error, LResult};
use engine::{Data, data_schema};

use crate::server::SqlServerImpl;

pub type QueryManagerRef = Arc<QueryManager>;

pub type QueryInfoRef = QueryInfo;

/// TODO queries will cause memory leak if not call get_data function
pub struct QueryManager {
    quries: DashMap<u64, QueryInfoRef>,
    id_generator: AtomicU64,
}

pub struct QueryInfo {
    id: u64,
    data: Data,
}

impl SqlServerImpl {
    pub async fn query(&self, sql: &str) -> LResult<(u64, SchemaRef)> {
        let data = self.db.query(sql).await?;
        let schema = data_schema(&data);
        let id = self.query_mgr.create_query_info(data);
        Ok((id, schema))
    }
}

impl QueryManager {
    pub fn default() -> Self {
        Self {
            quries: Default::default(),
            id_generator: AtomicU64::new(0),
        }
    }
    fn create_query_info(&self, data: Data) -> u64 {
        let id = self.id_generator.fetch_add(1, Ordering::SeqCst);
        let info = QueryInfo { id, data };
        self.quries.insert(id, info);
        id
    }

    pub async fn get_data(&self, id: u64) -> LResult<(SchemaRef, SendableRecordBatchStream)>
    {
        let info = self.quries.remove(&id).ok_or_else(|| {
            Error::ExecutionErr(format!("query id {} not exists", id))
        })?.1;
        let schema = data_schema(&info.data);
        let stream = match info.data {
            Data::Left(df) => {
                df.execute_stream().await.map_err(|e| {
                    Error::ExecutionErr(format!("execute err: {}", e))
                })?
            }
            Data::Right(r) => {
                let s = MemoryStream::try_new(vec![r], schema.clone(), None).map_err(|e| {
                    Error::ExecutionErr(format!("MemoryStream err: {}", e))
                })?;
                Box::pin(s)
            }
        };
        Ok((schema, stream))
    }
}